import os
import json
import time

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
import lightning as L
from lightning.fabric.fabric import _FabricOptimizer
from lightning.fabric.loggers import TensorBoardLogger, CSVLogger
from torchmetrics import MeanMetric

from model import BaseModel as Model
from dataset import load_datasets
from losses import FocalLoss, DiceLoss
from utils import calc_iou, AverageMeter
import segmentation_models_pytorch as smp
from logger import Logger


def validate(fabric: L.Fabric, model: Model, val_dataloader: DataLoader, epoch: int = 0):
    model.eval()
    ious = AverageMeter()
    f1_scores = AverageMeter()

    with torch.no_grad():
        for iter, data in enumerate(val_dataloader):
            images, bboxes, gt_masks = data
            num_images = images.size(0)
            pred_masks, _ = model(images, bboxes)
            for pred_mask, gt_mask in zip(pred_masks, gt_masks):
                batch_stats = smp.metrics.get_stats(
                    pred_mask,
                    gt_mask.int(),
                    mode='binary',
                    threshold=0.5,
                )
                batch_iou = smp.metrics.iou_score(*batch_stats, reduction="micro-imagewise")
                batch_f1 = smp.metrics.f1_score(*batch_stats, reduction="micro-imagewise")
                ious.update(batch_iou, num_images)
                f1_scores.update(batch_f1, num_images)
            fabric.print(
                f'Val: [{epoch}] - [{iter}/{len(val_dataloader)}]: Mean IoU: [{ious.avg:.4f}] -- Mean F1: [{f1_scores.avg:.4f}]'
            )

    fabric.print(f'Validation [{epoch}]: Mean IoU: [{ious.avg:.4f}] -- Mean F1: [{f1_scores.avg:.4f}]')
    fabric.loggers[0].log_metrics({
        'iou': ious.avg,
        'f1': f1_scores.avg,
    }, step=epoch)

    model.train()
    return f1_scores.avg


def train_sam(
        cfg: dict,
        fabric: L.Fabric,
        model: Model,
        optimizer: _FabricOptimizer,
        scheduler: _FabricOptimizer,
        train_dataloader: DataLoader,
        val_dataloader: DataLoader,
):
    """The SAM training loop."""
    focal_loss = FocalLoss()
    dice_loss = DiceLoss()
    f1_score = 0.0     
    best_f1 = 0.0 

    # 从中断处开始训练
    start_epoch = 1
    if cfg['resume']:
        map_location = 'cuda:%d' % fabric.global_rank
        checkpoint = torch.load(cfg['resume'], map_location={'cuda:0': map_location})
        start_epoch = checkpoint['epoch']
        network = checkpoint['network']
        opt = checkpoint['optimizer']
        sche = checkpoint['scheduler']
        model.model.load_state_dict(network)
        optimizer.load_state_dict(opt)
        scheduler.load_state_dict(sche)
        fabric.print(f'resume from {cfg["resume"]}')
    
    for epoch in range(start_epoch, cfg['num_epochs']):
        batch_time = AverageMeter(name="batch_time")
        data_time = AverageMeter(name="data_time")
        focal_losses = AverageMeter(name="focal_losses")
        dice_losses = AverageMeter(name="dice_losses")
        iou_losses = AverageMeter(name="iou_losses")
        total_losses = AverageMeter(name="total_losses")
        end = time.time()
        validated = False

        if not validated and epoch % cfg['val_interval'] == 0:
            f1_score = validate(fabric, model, val_dataloader, epoch)
            # 保存最好的验证模型
            if f1_score > best_f1:
                best_f1 = f1_score
                fabric.print(f'save best model to {cfg["out_dir"]}')
                state_dict = model.model.state_dict()
                if fabric.global_rank == 0:
                    torch.save(state_dict, os.path.join(cfg['out_dir'], f'best_{cfg["model"]["type"]}_f1{f1_score:.2f}.pth'))
            validated = True
        
        for iter, data in enumerate(train_dataloader):
            data_time.update(time.time() - end)
            images, bboxes, gt_masks = data
            batch_size = images.shape[0]
            pred_masks, iou_predictions = model(images, bboxes)
            num_masks = sum(len(pred_mask) for pred_mask in pred_masks)
            loss_focal = torch.tensor(0., device=fabric.device)
            loss_dice = torch.tensor(0., device=fabric.device)
            loss_iou = torch.tensor(0., device=fabric.device)
            for pred_mask, gt_mask, iou_prediction in zip(pred_masks, gt_masks, iou_predictions):
                batch_iou = calc_iou(pred_mask, gt_mask)
                loss_focal += focal_loss(pred_mask, gt_mask, num_masks)
                loss_dice += dice_loss(pred_mask, gt_mask, num_masks)
                loss_iou += F.mse_loss(iou_prediction, batch_iou, reduction='sum') / num_masks

            loss_total = 20. * loss_focal + loss_dice + loss_iou
            print(f'loss_total:{loss_total}')
            optimizer.zero_grad()
            fabric.backward(loss_total)
            optimizer.step()
            scheduler.step()
            batch_time.update(time.time() - end)
            end = time.time()

            focal_losses.update(loss_focal.item(), batch_size)
            dice_losses.update(loss_dice.item(), batch_size)
            iou_losses.update(loss_iou.item(), batch_size)
            total_losses.update(loss_total.item(), batch_size)

            fabric.print(f'Epoch: [{epoch}][{iter+1}/{len(train_dataloader)}]'
                         f' | Time [{batch_time.val:.3f}s ({batch_time.avg:.3f}s)]'
                         f' | Data [{data_time.val:.3f}s ({data_time.avg:.3f}s)]'
                         f' | Focal Loss [{focal_losses.val:.4f} ({focal_losses.avg:.4f})]'
                         f' | Dice Loss [{dice_losses.val:.4f} ({dice_losses.avg:.4f})]'
                         f' | IoU Loss [{iou_losses.val:.4f} ({iou_losses.avg:.4f})]'
                         f' | Total Loss [{total_losses.val:.4f} ({total_losses.avg:.4f})]')
        

        # 保存模型
        if epoch % cfg['save_interval'] == 0:
            fabric.print(f'save checkpoint to {cfg["out_dir"]}')
            state_dict = model.model.state_dict()
            # checkpoint = {
            #     'epoch': epoch,
            #     'network': state_dict,
            #     'optimizer': optimizer.state_dict(),
            #     'scheduler': scheduler.state_dict(),
            # }
            checkpoint = state_dict
            if fabric.global_rank == 0:
                torch.save(checkpoint, os.path.join(cfg['out_dir'], f'sam_{cfg["model"]["type"]}_{epoch:06d}-f1{f1_score:.2f}-ckpt.pth'))
        
        fabric.loggers[0].log_metrics({
            'focal_loss': focal_losses.avg,
            'dice_loss': dice_losses.avg,
            'iou_loss': iou_losses.avg,
            'total_loss': total_losses.avg,
        }, step=epoch)
            


def configure_opt(cfg: dict, model: Model):
    def lr_lambda(step):
        if step < cfg['opt']['warmup_steps']:
            return step / cfg['opt']['warmup_steps']
        elif step < cfg['opt']['steps'][0]:
            return 1.0
        elif step < cfg['opt']['steps'][1]:
            return 1 / cfg['opt']['decay_factor']
        else:
            return 1 / (cfg['opt']['decay_factor'] ** 2)
    
    optimizer = torch.optim.Adam(model.model.parameters(), lr=cfg['opt']['learning_rate'], weight_decay=cfg['opt']['weight_decay'])
    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
    return optimizer, scheduler


def main(cfg: dict) -> None:
    fabric = L.Fabric(accelerator="auto",
                   devices=cfg['num_devices'],
                   strategy="auto",
                   loggers=[TensorBoardLogger(cfg['out_dir'], name="lightning-sam")])
    fabric.launch()
    fabric.seed_everything(1337 + fabric.global_rank)

    if fabric.global_rank == 0:
        os.makedirs(cfg['out_dir'], exist_ok=True)

    # with fabric.device:
    #     model = Model(cfg)
    #     model.setup()
    device = torch.device(fabric.device)
    model = Model(cfg).to(device)
    model.setup()

    train_loder, val_loader = load_datasets(cfg, model.model.image_encoder.img_size)
    train_data = fabric._setup_dataloader(train_loder)
    val_data = fabric._setup_dataloader(val_loader)

    optimizer, scheduler = configure_opt(cfg, model)
    fabric.loggers[0].log_hyperparams(cfg)
    model, optimizer = fabric.setup(model, optimizer)

    train_sam(cfg, fabric, model, optimizer, scheduler, train_data, val_data)
    # validate(fabric, model, val_data, epoch=0)


if __name__ == "__main__":
    cfg = json.load(open('./configs/AAR.json'))
    main(cfg)